package simpletrace

import (
	"context"
	"encoding/json"
	"sort"
)

type SpeedscopeRenderer struct{}

func NewSpeedscopeRenderer() *SpeedscopeRenderer {
	return &SpeedscopeRenderer{}
}

func (r *SpeedscopeRenderer) Display(ctx context.Context, span *Span, opt any) (string, error) {
	type frameKey string
	frameMap := map[frameKey]int{}   // function name → index
	frameList := []speedscopeFrame{} // ordered list
	events := []speedscopeEvent{}    // flat events

	// 修复父子时间错乱：确保父包含子
	var fixTimestamps func(n *Span)
	fixTimestamps = func(n *Span) {
		if n == nil {
			return
		}
		for _, c := range n.Children {
			fixTimestamps(c)
		}
		for _, c := range n.Children {
			if c == nil {
				continue
			}
			if c.Start.Before(n.Start) {
				n.Start = c.Start
			}
			if c.End.After(n.End) {
				n.End = c.End
			}
		}
	}
	fixTimestamps(span)

	baseTime := span.Start

	// assign index to frames
	getFrameIndex := func(name string) int {
		key := frameKey(name)
		if idx, ok := frameMap[key]; ok {
			return idx
		}
		idx := len(frameList)
		frameMap[key] = idx
		frameList = append(frameList, speedscopeFrame{Name: name})
		return idx
	}

	// DFS遍历生成事件
	var dfs func(n *Span)
	dfs = func(n *Span) {
		if n == nil {
			return
		}
		idx := getFrameIndex(n.Name)
		start := n.Start.Sub(baseTime).Seconds() * 1000
		events = append(events, speedscopeEvent{
			Type:  "O",
			Frame: idx,
			At:    start,
		})
		for _, c := range n.Children {
			dfs(c)
		}
		if !n.End.IsZero() {
			end := n.End.Sub(baseTime).Seconds() * 1000
			events = append(events, speedscopeEvent{
				Type:  "C",
				Frame: idx,
				At:    end,
			})
		}
	}

	dfs(span)

	// 按时间排序事件
	sort.Slice(events, func(i, j int) bool {
		return events[i].At < events[j].At
	})

	// 构造 Speedscope JSON 文件
	var file speedscopeFile
	file.Schema = "https://www.speedscope.app/file-format-schema.json"
	file.Name = "TraceContext Flamegraph"
	file.Shared.Frames = frameList
	file.Profiles = []speedscopeProfile{
		{
			Type:       "evented",
			Name:       "Trace",
			Unit:       "milliseconds",
			Events:     events,
			StartValue: events[0].At,
			EndValue:   events[len(events)-1].At,
		},
	}

	res, err := json.MarshalIndent(file, "", "  ")
	return string(res), err
}

// speedscopeEvent represents an Open (O) or Close (C) event
type speedscopeEvent struct {
	Type  string  `json:"type"`  // "O" or "C"
	Frame int     `json:"frame"` // index into frames list
	At    float64 `json:"at"`    // timestamp in milliseconds
}

// speedscopeProfile is one profile block in speedscope format
type speedscopeProfile struct {
	Type       string            `json:"type"` // "evented"
	Name       string            `json:"name"`
	Unit       string            `json:"unit"` // "milliseconds"
	Events     []speedscopeEvent `json:"events"`
	StartValue float64           `json:"startValue"` // optional, start time in milliseconds
	EndValue   float64           `json:"endValue"`   // optional, end time in milliseconds
}

// speedscopeFrame holds shared frame metadata
type speedscopeFrame struct {
	Name string `json:"name"`
}

// speedscopeFile is the top-level structure
type speedscopeFile struct {
	Schema   string              `json:"$schema"`
	Name     string              `json:"name"`
	Profiles []speedscopeProfile `json:"profiles"`
	Shared   struct {
		Frames []speedscopeFrame `json:"frames"`
	} `json:"shared"`
}
